#ifndef KEYGEN_H
#define KEYGEN_H
#include <vector>
#include "ring.h"
#include "ring_poly.h"
#include "context.h"
#include "prng.h"
#include "modular_reduction.h"
#include "gmp.h"
#include "ring_sampler.h"
#include <unordered_map>

class PublicKey;
class SecretKey;
class SwitchingKey;
class RotationKeySet;
class RelinearizationKey;
class GaloisKeys;
class EvaluationKey;
/* *
 * @class KeyGenerator
 * @brief 密钥生成器类，用于生成同态加密所需的各类密钥。
 *
 * `KeyGenerator` 类提供了生成加密和解密所需的各种密钥的功能，包括秘密密钥、公钥、旋转密钥、Galois密钥等。
 */
class KeyGenerator {
public:
    Context *parameters;
    Poly *poolQ;
    PolyQP *poolQP;
    GaussianSampler *gaussianSamplerQ;
    UniformSampler *uniformSamplerQ;
    UniformSampler *uniformSamplerP;

    KeyGenerator(Context *parameters);
    // 拷贝构造函数
    KeyGenerator(const KeyGenerator &) = delete;
    KeyGenerator &operator = (const KeyGenerator &) = delete;
    // 移动构造函数
    KeyGenerator(KeyGenerator &&other) = delete;
    KeyGenerator &operator = (KeyGenerator &&other) = delete;

    std::pair<SecretKey *, PublicKey *> GenKeyPair();
    SecretKey *GenSecretKey();
    PublicKey *GenPublicKey(SecretKey *sk);
    SecretKey *GenSecretkeyWithDistrib(double p);
    SecretKey *genSecretKeyFromSampler(BaseSampler *sampler);
    RotationKeySet *GenRotationKeysForRotations(std::vector<int> ks, bool includeConjugate, SecretKey *sk);
    RotationKeySet *GenRotationKeys(std::vector<uint64_t> &galEls, SecretKey *sk);
    GaloisKeys *GenGaloisKeys(std::vector<uint64_t> &galEls, SecretKey *sk);
    void GenrotKey(PolyQP sk, uint64_t galEl, SwitchingKey *swk);
    RelinearizationKey *GenRelinearizationKey(SecretKey *sk, int maxDegree);
    RotationKeySet *GenRotationKeysForInnerSum(SecretKey *sk);
    SwitchingKey *GenSwitchingKey(SecretKey *skInput, SecretKey *skOutput);

private:
    void PrivateGenSwitchingKey(Poly *skIn, PolyQP *skOut, SwitchingKey *swk);
};
/* *
 * @class PublicKey
 * @brief 公钥。
 */
class PublicKey {
public:
    std::vector<PolyQP> value;

    PublicKey(Context *parameters);
};
/* *
 * @class SecretKey
 * @brief 私钥。
 */
class SecretKey {
public:
    PolyQP value;

    SecretKey(Context *parameters);
    // 复制构造函数
    SecretKey(const SecretKey &other);
    SecretKey &operator = (const SecretKey &other);
};
/* *
 * @class SwitchingKey
 * @brief 交换密钥。
 */
class SwitchingKey {
public:
    std::vector<std::vector<PolyQP *>> value;

    SwitchingKey(Context *parameters, int levelQ, int levelP);
};
/* *
 * @class RotationKeySet
 * @brief 旋转密钥。
 */
class RotationKeySet {
public:
    std::unordered_map<uint64_t, SwitchingKey *> keys;

    RotationKeySet(Context *parameters, const std::vector<uint64_t> &galoisElement);
    SwitchingKey *GetRotationKey(uint64_t galoisEl);
};
/* *
 * @class RelinearizationKey
 * @brief 重线性化密钥。
 */
class RelinearizationKey {
public:
    std::vector<SwitchingKey *> keys;
    RelinearizationKey() = default;

    RelinearizationKey(Context *parameters, int maxRelinDegree);
};
/* *
 * @class GaloisKeys
 * @brief 伽罗瓦密钥。
 */
class GaloisKeys {
public:
    std::unordered_map<uint64_t, SwitchingKey *> keys;
    GaloisKeys() = default;
    GaloisKeys(Context *parameters, const std::vector<uint64_t> &galoisElement);
    SwitchingKey *GetGaloisKey(uint64_t galoisEl);
};
/* *
 * @class EvaluationKey
 * @brief 评估密钥，包含重线性化密钥和旋转密钥。
 */
class EvaluationKey {
public:
    RelinearizationKey *rlk;
    RotationKeySet *rtks;

    EvaluationKey(RelinearizationKey *rlk, RotationKeySet *rtks);
};

#endif // KEYGEN_H